import networkx as nx
import numpy as np
from torch.utils import data
from torch.utils.data import DataLoader
import torch


def Read_graph(file_name):
    # 文本文件中的每一行必须含有相同的数据; delimiter分隔符默认是空格; 类型是numpy array
    edge = np.loadtxt(file_name).astype(np.int32)
    # 得到图中点的最小和最大编号; .min()返回数组中所有元素最小的
    min_node, max_node = edge.min(), edge.max()
    # Node表示图上一共有多少个顶点，如果标号是从0开始，那么顶点数 = max_node + 1
    if min_node == 0:
        Node = max_node + 1
    else:
        Node = max_node
    # 这里面使用networkx将图的信息存入
    G = nx.Graph()
    # Adj就是图的邻接表矩阵，是一个n*n大小的numpy矩阵，这里n是顶点的个数
    Adj = np.zeros([Node, Node], dtype=np.int32)
    # 遍历边的文件，将每条边存入networkx的图，以及邻接矩阵Adj所对应的位置(i, j)
    for i in range(edge.shape[0]):
        G.add_edge(edge[i][0], edge[i][1])
        if min_node == 0:
            Adj[edge[i][0], edge[i][1]] = 1
            Adj[edge[i][1], edge[i][0]] = 1
        else:
            Adj[edge[i][0] - 1, edge[i][1] - 1] = 1
            Adj[edge[i][1] - 1, edge[i][0] - 1] = 1
    # 转化成tensor
    Adj = torch.FloatTensor(Adj)
    return G, Adj, Node


# 继承自pytorch中的data.Dataset类，为了后面batch训练方便
class Dataload(data.Dataset):

    def __init__(self, Adj, Node):
        self.Adj = Adj
        self.Node = Node

    def __getitem__(self, index):
        return index
        # adj_batch = self.Adj[index]
        # adj_mat = adj_batch[index]
        # b_mat = torch.ones_like(adj_batch)
        # b_mat[adj_batch != 0] = self.Beta
        # return adj_batch, adj_mat, b_mat

    def __len__(self):
        return self.Node


if __name__ == '__main__':
    G, Adj, Node = Read_graph('./karate/karate.edgelist')
    Data = Dataload(Adj, Node)
    Test = DataLoader(Data, batch_size=20, shuffle=True)
    for index in Test:
        adj_batch = Adj[index]
        adj_mat = adj_batch[:, index]
        b_mat = torch.ones_like(adj_batch)
        b_mat[adj_batch != 0] = 5
        print(b_mat)
